﻿using Centny.PdfReader.Pdf;
using System;
using System.Collections.Generic;
using System.Linq;
using System.Text;
using System.IO;
using System.Threading.Tasks;
using Windows.Security.Cryptography;
using Windows.Security.Cryptography.Core;
using Windows.Storage.Streams;
using Centny.PdfReader.Types;

namespace Centny.PdfReader.Security
{
    public class PdfStandardSecurityHandler
    {
        public static byte[] DefaultPwd = 
        {
            0x28, 0xbf, 0x4e, 0x5e, 0x4e, 0x75, 0x8a, 0x41,//
            0x64, 0x00, 0x4e, 0x56, 0xff, 0xfa, 0x01, 0x08,//
            0x2e, 0x2e, 0x00, 0xb6, 0xd0, 0x68, 0x3e, 0x80,//
            0x2f, 0x0c, 0xa9, 0xfe, 0x64, 0x53, 0x69, 0x7a//
        };
        public PdfClassObject Encrypt { get; set; }
        public PdfDictionary Trailer { get; set; }
        byte[] DId { get; set; }
        byte[] OValue { get; set; }
        byte[] UValue { get; set; }
        int PValue { get; set; }
        int RValue { get; set; }
        //public int RValue;
        public byte[] GlobalKey { get; set; }
        //
        public PdfStandardSecurityHandler(PdfClassObject encrypt, PdfDictionary trailer)
        {
            this.Encrypt = encrypt;
            this.Trailer = trailer;
            //document ID.
            PdfString d = (this.Trailer[TrailerKeys.ID] as PdfArray)[0] as PdfString;
            this.DId = d.Bys;
            //encrypt /O.
            PdfString o = this.Encrypt.AsDictionary[EncryptKeys.O] as PdfString;
            this.OValue = o.Bys;
            //encrypt /U.
            PdfString u = this.Encrypt.AsDictionary[EncryptKeys.U] as PdfString;
            this.UValue = u.Bys;
            //encrypt /P
            PdfLong p = this.Encrypt.AsDictionary[EncryptKeys.P] as PdfLong;
            this.PValue = (int)p.Val;
            //encrypt /R
            PdfLong r = this.Encrypt.AsDictionary[EncryptKeys.R] as PdfLong;
            this.RValue = (int)r.Val;
        }
        protected byte[] GenerateEncyptPwd(byte[] oitem, String pwd)
        {
            byte[] nitem = new byte[oitem.Length];
            Array.Copy(oitem, nitem, oitem.Length);
            byte[] ppwd = PadPassword(pwd);
            byte[] digest = Md5(ppwd);
            MemoryStream ms = new MemoryStream();
            if (this.RValue > 2)
            {
                for (int i = 0; i < 50; i++)
                {
                    digest = Md5(digest);
                }
                byte[] mkey = new byte[16];
                for (int i = 0; i < 20; i++)
                {
                    for (int j = 0; j < mkey.Length; ++j)
                        mkey[j] = (byte)(digest[j] ^ i);
                    //this.PrepareRC4Key(mkey, 0, mkey.Length);
                    //this.EncryptRC4(nitem, 0, nitem.Length, nitem);
                    nitem = Rc4Decryption(nitem, mkey);
                }
            }
            else
            {
                nitem = Rc4Decryption(nitem, digest, 5);
            }
            return nitem;
        }
        protected byte[] PadPassword(String pwd)
        {
            if (pwd == null)
            {
                pwd = "";
            }
            byte[] bys = new byte[DefaultPwd.Length];
            byte[] pby = Encoding.UTF8.GetBytes(pwd);
            if (pby.Length > DefaultPwd.Length)
            {
                Array.Copy(pby, bys, DefaultPwd.Length);
            }
            else
            {
                Array.Copy(pby, bys, pby.Length);
                Array.Copy(DefaultPwd, 0, bys, pby.Length, DefaultPwd.Length - pby.Length);
            }
            return bys;
        }
        public bool GenerateGlobalKey(String pwd, bool IsOwnerPwd = false)
        {
            MemoryStream buf = new MemoryStream();
            byte[] rpwd;
            if (IsOwnerPwd)
            {
                rpwd = this.GenerateEncyptPwd(this.OValue, pwd);
            }
            else
            {
                rpwd = this.PadPassword(pwd);
            }
            buf.Write(rpwd, 0, rpwd.Length);
            buf.Write(this.OValue, 0, this.OValue.Length);
            byte[] pbys = BitConverter.GetBytes((UInt32)this.PValue);
            buf.Write(pbys, 0, 4);
            buf.Write(this.DId, 0, this.DId.Length);
            //
            //
            //create MD5.
            byte[] bys = buf.ToArray();
            byte[] md5 = Md5(bys);
            if (this.RValue > 2)
            {
                for (int idx = 0; idx < 50; idx++)
                {
                    md5 = Md5(md5);
                }
                this.GlobalKey = new byte[16];
            }
            else
            {
                this.GlobalKey = new byte[5];
            }
            Array.Copy(md5, this.GlobalKey, this.GlobalKey.Length);
            this.GenerateUserKey();
            return this.EqualUserKey();
        }
        protected byte[] UserKey = new byte[32];
        protected byte[] GenerateUserKey()
        {
            if (this.RValue > 2)
            {
                MemoryStream ms = new MemoryStream();
                ms.Write(DefaultPwd, 0, DefaultPwd.Length);
                ms.Write(this.DId, 0, this.DId.Length);
                byte[] digest = Md5(ms.ToArray());
                byte[] bkey = new byte[16];
                Array.Copy(digest, 0, bkey, 0, 16);
                for (int i = 0; i < 20; i++)
                {
                    for (int j = 0; j < this.GlobalKey.Length; j++)
                        digest[j] = (byte)(this.GlobalKey[j] ^ i);
                    Rc4Decryption(bkey, bkey, digest, digest.Length);
                }
                Array.Copy(bkey, this.UserKey, 16);
                for (int idx = 16; idx < 32; idx++)
                    this.UserKey[idx] = 0;
            }
            else
            {
                Rc4Decryption(DefaultPwd, this.UserKey, this.GlobalKey, this.GlobalKey.Length);
            }
            return this.UserKey;
        }
        public bool EqualUserKey()
        {
            for (int i = 0; i < 32; i++)
            {
                if (this.UserKey[i] != this.UValue[i])
                {
                    return false;
                }
            }
            return true;
        }
        public byte[] GenerateObjectKey(PdfObject obj)
        {
            MemoryStream buf = new MemoryStream();
            //global key.
            buf.Write(this.GlobalKey, 0, this.GlobalKey.Length);
            //object number.
            byte[] numbys = BitConverter.GetBytes(obj.Num);
            buf.Write(numbys, 0, 3);
            //object modify.
            byte[] mtime = BitConverter.GetBytes(obj.ModifyTimes);
            buf.Write(mtime, 0, 2);
            //
            //
            //create md5.
            byte[] md5 = Md5(buf.ToArray());
            int ksize = this.GlobalKey.Length + 5;
            if (ksize > 16)
            {
                ksize = 16;
            }
            byte[] okey = new byte[ksize];
            Array.Copy(md5, okey, ksize);
            return okey;
        }
        public static byte[] Md5(byte[] src)
        {
            using (var dw = new DataWriter())
            {
                dw.WriteBytes(src);
                HashAlgorithmProvider objAlgProv = HashAlgorithmProvider.OpenAlgorithm(HashAlgorithmNames.Md5);
                // Hash the message.
                IBuffer buffHash = objAlgProv.HashData(dw.DetachBuffer());
                if (buffHash.Length != objAlgProv.HashLength)
                {
                    throw new Exception("There was an error creating the Md5 hash");
                }
                byte[] res = new byte[buffHash.Length];
                var dr = DataReader.FromBuffer(buffHash);
                dr.ReadBytes(res);
                return res;
            }
        }
        /// <summary>
        /// Bytes used for RC4 encryption.
        /// </summary>
        byte[] state = new byte[256];
        /// <summary>
        /// Prepare the encryption key.
        /// </summary>
        void PrepareRC4Key(byte[] key, int offset, int length)
        {
            int idx1 = 0;
            int idx2 = 0;
            for (int idx = 0; idx < 256; idx++)
                this.state[idx] = (byte)idx;
            byte tmp;
            for (int idx = 0; idx < 256; idx++)
            {
                idx2 = (key[idx1 + offset] + this.state[idx] + idx2) & 255;
                tmp = this.state[idx];
                this.state[idx] = this.state[idx2];
                this.state[idx2] = tmp;
                idx1 = (idx1 + 1) % length;
            }
        }
        /// <summary>
        /// Encrypts the data.
        /// </summary>
        void EncryptRC4(byte[] inputData, int offset, int length, byte[] outputData)
        {
            length += offset;
            int x = 0, y = 0;
            byte b;
            for (int idx = offset; idx < length; idx++)
            {
                x = (x + 1) & 255;
                y = (this.state[x] + y) & 255;
                b = this.state[x];
                this.state[x] = this.state[y];
                this.state[y] = b;
                outputData[idx] = (byte)(inputData[idx] ^ state[(this.state[x] + this.state[y]) & 255]);
            }
        }
        public static byte[] Rc4Decryption(byte[] src, byte[] key, int klen = int.MaxValue)
        {
            byte[] dst = new byte[src.Length];
            Rc4Decryption(src, dst, key, klen);
            return dst;
        }
        public static void Rc4Decryption(byte[] src, byte[] dst, byte[] key, int klen = int.MaxValue)
        {
            IBuffer kbuf;
            using (var dw = new DataWriter())
            {
                if (klen == int.MaxValue)
                {
                    dw.WriteBytes(key);
                }
                else
                {
                    byte[] nkey = new byte[klen];
                    Array.Copy(key, nkey, klen);
                    dw.WriteBytes(nkey);
                }
                kbuf = dw.DetachBuffer();
            }
            IBuffer sbuf;
            using (var dw = new DataWriter())
            {
                dw.WriteBytes(src);
                sbuf = dw.DetachBuffer();
            }
            IBuffer buffDecrypted;
            SymmetricKeyAlgorithmProvider objAlg = SymmetricKeyAlgorithmProvider.OpenAlgorithm(SymmetricAlgorithmNames.Rc4);
            CryptographicKey ckey = objAlg.CreateSymmetricKey(kbuf);
            buffDecrypted = CryptographicEngine.Decrypt(ckey, sbuf, null);
            using (var dr = DataReader.FromBuffer(buffDecrypted))
            {
                dr.ReadBytes(dst);
            }
        }
        sealed partial class SymmKeyAlgProviderApp
        {
            public SymmKeyAlgProviderApp()
            {

                // Initialize the encryption process.
                String strMsg = "1234567812345678";     // Data to encrypt.
                String strAlgName = SymmetricAlgorithmNames.AesCbc;
                UInt32 keyLength = 32;                  // Length of the key, in bytes
                BinaryStringEncoding encoding;          // Binary encoding value
                IBuffer iv;                             // Initialization vector
                CryptographicKey key;                   // Symmetric key

                // Encrypt a message.
                IBuffer buffEncrypted = this.SampleCipherEncryption(
                    strMsg,
                    strAlgName,
                    keyLength,
                    out encoding,
                    out iv,
                    out key);

                // Decrypt a message.
                this.SampleCipherDecryption(
                    strAlgName,
                    buffEncrypted,
                    iv,
                    encoding,
                    key);
            }

            public IBuffer SampleCipherEncryption(
                String strMsg,
                String strAlgName,
                UInt32 keyLength,
                out BinaryStringEncoding encoding,
                out IBuffer iv,
                out CryptographicKey key)
            {
                // Initialize the initialization vector.
                iv = null;

                // Initialize the binary encoding value.
                encoding = BinaryStringEncoding.Utf8;

                // Create a buffer that contains the encoded message to be encrypted. 
                IBuffer buffMsg = CryptographicBuffer.ConvertStringToBinary(strMsg, encoding);

                // Open a symmetric algorithm provider for the specified algorithm. 
                SymmetricKeyAlgorithmProvider objAlg = SymmetricKeyAlgorithmProvider.OpenAlgorithm(strAlgName);

                // Demonstrate how to retrieve the name of the algorithm used.
                String strAlgNameUsed = objAlg.AlgorithmName;

                // Determine whether the message length is a multiple of the block length.
                // This is not necessary for PKCS #7 algorithms which automatically pad the
                // message to an appropriate length.
                if (!strAlgName.Contains("PKCS7"))
                {
                    if ((buffMsg.Length % objAlg.BlockLength) != 0)
                    {
                        throw new Exception("Message buffer length must be multiple of block length.");
                    }
                }

                // Create a symmetric key.
                IBuffer keyMaterial = CryptographicBuffer.GenerateRandom(keyLength);
                key = objAlg.CreateSymmetricKey(keyMaterial);

                // CBC algorithms require an initialization vector. Here, a random
                // number is used for the vector.
                if (strAlgName.Contains("CBC"))
                {
                    iv = CryptographicBuffer.GenerateRandom(objAlg.BlockLength);
                }

                // Encrypt the data and return.
                IBuffer buffEncrypt = CryptographicEngine.Encrypt(key, buffMsg, iv);
                return buffEncrypt;
            }

            public void SampleCipherDecryption(
                String strAlgName,
                IBuffer buffEncrypt,
                IBuffer iv,
                BinaryStringEncoding encoding,
                CryptographicKey key)
            {
                // Declare a buffer to contain the decrypted data.
                IBuffer buffDecrypted;

                // Open an symmetric algorithm provider for the specified algorithm. 
                SymmetricKeyAlgorithmProvider objAlg = SymmetricKeyAlgorithmProvider.OpenAlgorithm(strAlgName);

                // The input key must be securely shared between the sender of the encrypted message
                // and the recipient. The initialization vector must also be shared but does not
                // need to be shared in a secure manner. If the sender encodes a message string 
                // to a buffer, the binary encoding method must also be shared with the recipient.
                buffDecrypted = CryptographicEngine.Decrypt(key, buffEncrypt, iv);

                // Convert the decrypted buffer to a string (for display). If the sender created the
                // original message buffer from a string, the sender must tell the recipient what 
                // BinaryStringEncoding value was used. Here, BinaryStringEncoding.Utf8 is used to
                // convert the message to a buffer before encryption and to convert the decrypted
                // buffer back to the original plaintext.
                String strDecrypted = CryptographicBuffer.ConvertBinaryToString(encoding, buffDecrypted);
            }
        }
    }
}
